'''
Version: 1.0
Author: hpf
Date: 2021-06-08 03:05:28
LastEditTime: 2021-07-29 12:53:32
LastEditors: hpf
Description: 跟踪算法
FilePath: /trackrunner/trackrunner.py
'''

import importlib
import os
import os.path as osp
import shutil
import sys

import cv2
import numpy as np

import FrameWork.datasets as track_datasets
import FrameWork.evaluation as track_evaluator
import FrameWork.utils as track_utils
'''
@description: 视频处理接口
@param {*} library_path 运行的算法路径
@param {*} model_path 运行算法的模型路径
@param {*} config_path 运行算法的配置文件路径
@param {*} input_path 输入算法处理的视频路径
@param {*} output_path 输出结果的文件夹
@return {*}
'''


def run_video(library_path, model_path, config_path, input_path, output_path):
    print('Starting run video...lib:{},model:{},config:{},input:{},output:{}'.
          format(library_path, model_path, config_path, input_path,
                 output_path))

    process_path = sys.argv[0]
    sys.argv.clear()
    sys.argv.append(process_path)

    # Import algorithm library
    os.chdir(library_path)
    sys.path.append(library_path)
    track = importlib.import_module("Trackor.track")

    # 加载配置文件
    if config_path == "":
        config_path = library_path + "/track_cfg.yaml"
    cfg_yaml = track_utils.YamlParser(config_file=config_path)

    if model_path != "":
        if cfg_yaml.EVAL.get('WEIGHT'):
            cfg_yaml.EVAL.WEIGHT = model_path

    # Load configuration file
    track_utils.mkdir_if_missing(output_path)
    timer = track_utils.Timer()
    dataloader = track_datasets.LoadVideo(input_path, preprocess=False)
    video_count = dataloader.vn
    process_num = -1
    results = []
    frame_id = 0

    # Initialize the tracking algorithm
    tracker = track.ZJLAB_TRACK(cfg_yaml.EVAL)

    # 遍历图片
    for _, img in dataloader:
        # Iterate over pictures
        if int(100 * frame_id / video_count) != process_num:
            process_num = int(100 * frame_id / video_count)
            print("PlatformProcessNum:%d" % (process_num))

        # Run to process pictures
        timer.tic()
        try:
            online_targets = tracker.run(img)
        except:
            print("error")
        timer.toc()
        if len(online_targets) > 0:
            online_tlwhs = online_targets[:, :4]
            online_ids = online_targets[:, -1]
            # Save result
            results.append((frame_id + 1, online_tlwhs, online_ids))
            online_im = track_utils.plot_tracking(img,
                                                  online_tlwhs,
                                                  online_ids,
                                                  frame_id=frame_id,
                                                  fps=1. / timer.average_time)
            cv2.imwrite(os.path.join(output_path, "%010d.jpg" % (frame_id)),
                        online_im)
        frame_id += 1
    print("PlatformProcessNum:%d" % (100))


'''
@description: 数据集处理接口
@param {*} library_path 运行的算法路径
@param {*} model_path 运行算法的模型路径
@param {*} config_path 运行算法的配置文件路径
@param {*} input_path 输入算法处理的数据集路径
@param {*} output_path 输出结果的文件夹
@return {*}
'''


def run_eval(library_path, model_path, config_path, input_path, output_path):

    print('Starting run eval...lib:{},model:{},config:{},input:{},output:{}'.
          format(library_path, model_path, config_path, input_path,
                 output_path))

    process_path = sys.argv[0]
    sys.argv.clear()
    sys.argv.append(process_path)

    # Import algorithm library
    os.chdir(library_path)
    sys.path.append(library_path)
    track = importlib.import_module("Trackor.track")
    # Load configuration file
    input_path = input_path + "/images/train"
    if config_path == "":
        config_path = library_path + "/track_cfg.yaml"
    cfg_yaml = track_utils.YamlParser(config_file=config_path)

    if model_path != "":
        if cfg_yaml.EVAL.get('WEIGHT'):
            cfg_yaml.EVAL.WEIGHT = model_path

    # Initialize variables
    seqs_str = '''MOT17-02-SDP
                      MOT17-04-SDP
                      MOT17-05-SDP
                      MOT17-09-SDP
                      MOT17-10-SDP
                      MOT17-11-SDP
                      MOT17-13-SDP
                    '''
    seqs = [seq.strip() for seq in seqs_str.split()]
    track_utils.mkdir_if_missing(output_path)
    save_txt = os.path.join(output_path, 'txt')
    track_utils.mkdir_if_missing(save_txt)
    data_type = 'mot'
    accs = []
    n_frame = 0
    timer_avgs, timer_calls = [], []
    seq_id = 0

    # Initialize the tracking algorithm
    tracker = track.ZJLAB_TRACK(cfg_yaml.EVAL)

    # Iterating over the data set
    for seq in seqs:
        print('start seq: %s' % (seq))
        dataloader = track_datasets.LoadImages(osp.join(
            input_path, seq, 'img1'),
                                               preprocess=False)
        timer = track_utils.Timer()
        results = []
        frame_id = 0
        for img0 in dataloader:
            # Run to process pictures
            timer.tic()
            try:
                online_targets = tracker.run(img0)
            except:
                print("error")
            
            timer.toc()
            if len(online_targets) > 0:
                online_tlwhs = online_targets[:, :4]
                online_ids = online_targets[:, -1]
                results.append((frame_id + 1, online_tlwhs, online_ids))
                online_im = track_utils.plot_tracking(img0,
                                                      online_tlwhs,
                                                      online_ids,
                                                      frame_id=frame_id,
                                                      fps=1. /
                                                      timer.average_time)
                cv2.imwrite(
                    os.path.join(output_path,
                                 seq + '{:08d}.jpg'.format(frame_id)),
                    online_im)
            frame_id += 1
        print("PlatformProcessNum:%d" % (100 * seq_id / len(seqs)))
        seq_id += 1

        # Save result
        result_filename = os.path.join(save_txt, '{}.txt'.format(seq))
        track_utils.write_results(result_filename, results, data_type)
        n_frame += frame_id
        timer_avgs.append(timer.average_time)
        timer_calls.append(timer.calls)

        # statistics
        evaluator = track_evaluator.Evaluator(input_path, seq, data_type)
        accs.append(evaluator.eval_file(result_filename))

    # Delete txt directory
    shutil.rmtree(save_txt)
    timer_avgs = np.asarray(timer_avgs)
    timer_calls = np.asarray(timer_calls)
    all_time = np.dot(timer_avgs, timer_calls)
    avg_time = all_time / np.sum(timer_calls)

    # Get summary
    summary, strsummary = track_evaluator.Evaluator.get_summary(accs,
                                                                seqs,
                                                                metrics=None)
    summary.insert(0, 'FPS', 1.0 / avg_time, allow_duplicates=False)

    print("PlatformSummary:%s" % (summary.to_json()))
    print("PlatformProcessNum:100")


if __name__ == '__main__':
    if int(sys.argv[6]) == 1:
        run_eval(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4],
                 sys.argv[5])
    else:
        run_video(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4],
                  sys.argv[5])
